#pragma once

#include <string>
#include <vector>
#include <map>
#include<set>

#include <google/protobuf/text_format.h>
#include <google/protobuf/io/zero_copy_stream_impl.h>
#include "onnx_pb.h"


namespace vkinfer
{
    struct OnnxWeight
    {
        std::string name;
        std::vector<uint32_t> shape;
        void* values = nullptr;

        uint32_t count() const 
        {
            if (shape.size() == 0) return 0;
            uint32_t acc = 1;
            for (auto& s : shape) acc *= s;
            return acc;
        }
    };

    enum class OnnxAttribValueType
    {
        VALUE_TYPE_FLOAT = 0,
        VALUE_TYPE_INT = 1,
        VALUE_TYPE_FLOATS = 2,
        VALUE_TYPE_INTS = 3,
    };

    struct OnnxAttribute
    {
        std::string name;
        OnnxAttribValueType type;
        union {
            float f32[8];
            int32_t i32[8];
        } values;
        uint32_t num_values;

        float get_float(uint32_t index) const;
        uint32_t get_int(uint32_t index) const;
    };

    struct OnnxNodeInfo
    {
        std::string name;
        std::string op_type;
        std::vector<std::string> weight_names;
        std::vector<std::string> input_names;
        std::vector<std::string> output_names;
        std::vector<OnnxAttribute> attributes;
    };

    class OnnxParser
    {
    private:
        onnx::GraphProto graph;
        std::set<std::string> weight_names;

    public:
        OnnxParser();
        ~OnnxParser() {}
        bool open(const std::string& model_file);

        void load_io_info(std::vector<std::string>& input_names, std::vector<std::string>& output_names);
        void load_weights(std::map<std::string, std::shared_ptr<OnnxWeight>>& weights);
        void load_network(std::map<std::string, std::shared_ptr<OnnxNodeInfo>>& nodes);
    };
}
